Sticking the landing: A simple reduced-variance gradient for ADVI

Compared to the REINFORCE gradient estimator, the reparameterization trick usually gives lower-variance estimators. We propose a simple variant of the standard reparameterized gradient estimator for the evidence lower bound that has even lower variance under certain circumstances. Specifically, we decompose the derivative with respect to the variational parameters into two parts: a path derivative and the score function. Removing the second term produces an unbiased gradient estimator whose variance approaches zero as the approximate posterior approaches the exact posterior. We propose that the removed term has arbitrarily high variance when the variational posterior has a complex form, as when using adaptive posteriors such as given by normalizing flows or stochastic Hamiltonian inference. 1 Estimators of the Evidence Lower Bound Variational inference posits a family of distributions Q and attempts to find an approximate posterior qφ by optimizing the evidence lower bound (ELBO): L(φ) = Ez∼q[log p(x, z)− log qφ(z |x)] (ELBO) An unbiased approximation of the gradient of the ELBO allows stochastic gradient descent to scalably learn complex models. When the joint distribution p(x, z) can be evaluated by p(x|z) and p(z) separately, the ELBO can be written in the following three forms: L(φ) = Ez∼q[log p(x|z) + log p(z)− log qφ(z|x)] (1) = Ez∼q[log p(x|z) + log p(z))] +H[qφ] (2) = Ez∼q[log p(x|z)]−KL(qφ(z|x)||p(z)) (3) By sampling z ∼ q(z), equations (1), (2), and (3) can be used to construct Monte-Carlo estimates of the ELBO. When p(z) and q(z|x) are multivariate Gaussians, using (3) is appealing because it analytically integrates out terms that would otherwise have to be estimated by Monte Carlo. Intuitively, we might that using exact integrals wherever possible will give lower-variance estimators. Surprisingly, though, there are circumstances under which (1), which we call the full Monte Carlo estimator of the ELBO, has lower variance than the estimator based on (3) which calculates the KL divergence exactly. Specifically, when q(z|x) = p(z|x), i.e. the variational approximation is exact, then the variance of the full Monte Carlo estimator is exactly zero, since its value is a constant